package com.li.flink.datastream.state

import org.apache.flink.api.common.functions.RichFlatMapFunction
import org.apache.flink.api.common.state.{MapState, MapStateDescriptor, ValueState, ValueStateDescriptor}
import org.apache.flink.api.scala._
import org.apache.flink.configuration.Configuration
import org.apache.flink.streaming.api.scala.StreamExecutionEnvironment
import org.apache.flink.util.Collector

import java.util.UUID

object StateApp {
  def main(args: Array[String]): Unit = {
    //获取运行环境
    val env = StreamExecutionEnvironment.getExecutionEnvironment
    test01(env)
    //test02(env)
    env.execute("StateApp")
  }

  /**
   * 使用valueState 求平均数
   *
   * @param env
   */
  def test01(env: StreamExecutionEnvironment): Unit = {

    env.fromCollection(List(
      (1L, 3L),
      (1L, 7L),
      (2L, 4L),
      (1L, 5L),
      (2L, 2L),
      (2L, 5L)
    )).keyBy(_._1)
      //      .flatMap(new CountWindowAverage())
      .flatMap(new CountWindowMapStateAverage())
      .print()

  }
}

class CountWindowMapStateAverage extends RichFlatMapFunction[(Long, Long), (Long, Double)] {

  private var sum: MapState[String, Long] = _

  override def flatMap(input: (Long, Long), out: Collector[(Long, Double)]): Unit = {
    sum.put(UUID.randomUUID().toString, input._2)
    //access the state value
    val it = sum.values().iterator()
    var i = 0
    var total = 0D
    while (it.hasNext) {
      i = i + 1
      total = total + it.next()
    }

    if (i == 3) {
      out.collect((input._1, total / 3))
      sum.clear()
    }
  }

  override def open(parameters: Configuration): Unit = {
    val map = new MapStateDescriptor[String, Long]("average", classOf[String], classOf[Long])
    sum = getRuntimeContext.getMapState(map)
  }
}

class CountWindowAverage extends RichFlatMapFunction[(Long, Long), (Long, Double)] {

  private var sum: ValueState[(Long, Double)] = _

  override def flatMap(input: (Long, Long), out: Collector[(Long, Double)]): Unit = {

    // access the state value
    val tmpCurrentSum = sum.value

    // If it hasn't been used before, it will be null
    val currentSum = if (tmpCurrentSum != null) {
      tmpCurrentSum
    } else {
      (0L, 0D)
    }

    // update the count
    val newSum = (currentSum._1 + 1, currentSum._2 + input._2)

    // update the state
    sum.update(newSum)

    // if the count reaches 2, emit the average and clear the state
    if (newSum._1 >= 3) {
      out.collect((input._1, newSum._2 / newSum._1))
      sum.clear()
    }
  }

  override def open(parameters: Configuration): Unit = {
    sum = getRuntimeContext.getState(
      new ValueStateDescriptor[(Long, Double)]("average", createTypeInformation[(Long, Double)])
    )
  }
}

